#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import json
import cv2
import numpy as np
from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS
import base64
from io import BytesIO
from PIL import Image
import pytesseract

# 初始化Flask应用
app = Flask(__name__)
CORS(app)  # 允许跨域请求

# 配置文件路径
CONFIG_FILE = 'config.json'

# 默认配置
DEFAULT_CONFIG = {
    "ocrEngine": "opencv",
    "tesseract": {
        "cmd": ""
    }
}

def load_config():
    """加载配置文件"""
    if os.path.exists(CONFIG_FILE):
        try:
            with open(CONFIG_FILE, 'r', encoding='utf-8') as f:
                config = json.load(f)
                # 合并默认配置和文件配置
                merged_config = DEFAULT_CONFIG.copy()
                merged_config.update(config)
                return merged_config
        except Exception as e:
            print(f"加载配置文件出错: {e}")
            return DEFAULT_CONFIG
    else:
        # 如果配置文件不存在，创建默认配置文件
        save_config(DEFAULT_CONFIG)
        return DEFAULT_CONFIG

def save_config(config):
    """保存配置到文件"""
    try:
        with open(CONFIG_FILE, 'w', encoding='utf-8') as f:
            json.dump(config, f, ensure_ascii=False, indent=2)
        return True
    except Exception as e:
        print(f"保存配置文件出错: {e}")
        return False

# 加载配置
app.config['APP_CONFIG'] = load_config()

@app.route('/')
def index():
    """主页路由"""
    return send_from_directory('.', 'index.html')

@app.route('/settings.html')
def settings():
    """设置页面路由"""
    return send_from_directory('.', 'settings.html')

@app.route('/<path:path>')
def serve_static(path):
    """静态文件服务"""
    return send_from_directory('.', path)

@app.route('/api/config', methods=['GET'])
def get_config():
    """获取当前配置"""
    config = app.config['APP_CONFIG']
    return jsonify(config)

@app.route('/api/config', methods=['POST'])
def save_config_api():
    """保存配置"""
    try:
        data = request.get_json()
        if not data:
            return jsonify({'success': False, 'error': '没有接收到数据'}), 400
        
        # 验证必需的字段
        if 'ocrEngine' not in data:
            return jsonify({'success': False, 'error': '缺少ocrEngine字段'}), 400
        
        # 更新配置
        config = app.config['APP_CONFIG']
        config['ocrEngine'] = data['ocrEngine']
        
        if 'tesseract' in data:
            config['tesseract'] = data['tesseract']
        
        # 保存配置
        if save_config(config):
            # 更新应用配置
            app.config['APP_CONFIG'] = config
            return jsonify({'success': True, 'message': '设置已保存', 'config': config})
        else:
            return jsonify({'success': False, 'error': '无法保存配置文件'}), 500
            
    except Exception as e:
        return jsonify({'success': False, 'error': str(e)}), 500

def preprocess_image_for_ocr(image):
    """预处理图像以提高OCR准确性"""
    # 转换为灰度图
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image.copy()
    
    # 应用高斯模糊以减少噪声
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # 应用自适应阈值处理
    thresh = cv2.adaptiveThreshold(
        blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2
    )
    
    # 形态学操作，清理图像
    kernel = np.ones((2,2), np.uint8)
    cleaned = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kernel)
    
    return cleaned

def extract_sudoku_grid(image):
    """提取数独网格区域"""
    # 预处理图像
    processed = preprocess_image_for_ocr(image)
    
    # 查找轮廓
    contours, _ = cv2.findContours(processed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # 寻找最大的四边形轮廓（假设为数独网格）
    largest_contour = None
    max_area = 0
    
    for contour in contours:
        # 计算轮廓面积
        area = cv2.contourArea(contour)
        if area > max_area:
            # 近似轮廓为多边形
            epsilon = 0.02 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            # 如果是四边形
            if len(approx) == 4:
                max_area = area
                largest_contour = approx
    
    if largest_contour is not None:
        # 透视变换提取网格
        # 重新排列角点顺序
        points = largest_contour.reshape(4, 2)
        # 计算四个角点
        s = points.sum(axis=1)
        diff = np.diff(points, axis=1)
        tl = points[np.argmin(s)]      # top-left
        br = points[np.argmax(s)]      # bottom-right
        tr = points[np.argmin(diff)]   # top-right
        bl = points[np.argmax(diff)]   # bottom-left
        
        # 定义目标尺寸
        size = 450
        pts1 = np.float32([tl, tr, bl, br])
        pts2 = np.float32([[0, 0], [size, 0], [0, size], [size, size]])
        
        # 应用透视变换
        matrix = cv2.getPerspectiveTransform(pts1, pts2)
        warped = cv2.warpPerspective(image, matrix, (size, size))
        
        return warped
    
    # 如果没有找到四边形，返回原图
    return image

def split_grid_to_cells(image):
    """将数独网格分割为81个单元格"""
    height, width = image.shape[:2]
    cell_height = height // 9
    cell_width = width // 9
    
    cells = []
    for i in range(9):
        row = []
        for j in range(9):
            # 提取单元格
            y1 = i * cell_height
            y2 = (i + 1) * cell_height
            x1 = j * cell_width
            x2 = (j + 1) * cell_width
            
            cell = image[y1:y2, x1:x2]
            row.append(cell)
        cells.append(row)
    
    return cells

def recognize_digit_with_opencv(cell):
    """使用OpenCV识别单个单元格中的数字"""
    # 预处理单元格
    processed = preprocess_image_for_ocr(cell)
    
    # 计算白色像素比例
    white_pixels = cv2.countNonZero(processed)
    total_pixels = processed.shape[0] * processed.shape[1]
    ratio = white_pixels / total_pixels if total_pixels > 0 else 0
    
    # 如果白色像素比例过小，认为是空单元格
    if ratio < 0.05:
        return 0
    
    # 这里可以添加更复杂的数字识别逻辑
    # 目前返回0表示未识别
    return 0

def recognize_digit_with_tesseract(cell, config):
    """使用Tesseract识别单个单元格中的数字"""
    try:
        # 预处理单元格
        processed = preprocess_image_for_ocr(cell)
        
        # 如果配置了Tesseract命令路径，设置它
        if config.get('tesseract', {}).get('cmd'):
            pytesseract.pytesseract.tesseract_cmd = config['tesseract']['cmd']
        
        # 使用pytesseract识别
        custom_config = r'--oem 3 --psm 10 -c tessedit_char_whitelist=123456789'
        text = pytesseract.image_to_string(processed, config=custom_config)
        
        # 提取数字
        digits = [char for char in text if char.isdigit()]
        if digits:
            return int(digits[0])
        return 0
    except Exception as e:
        print(f"Tesseract识别出错: {e}")
        return 0

@app.route('/api/ocr', methods=['POST'])
def ocr_process():
    """处理OCR请求"""
    try:
        data = request.get_json()
        if not data or 'image' not in data:
            return jsonify({'success': False, 'error': '缺少图像数据'}), 400
        
        # 解码base64图像数据
        image_data = data['image']
        if image_data.startswith('data:image'):
            image_data = image_data.split(',')[1]
        
        image_bytes = base64.b64decode(image_data)
        image = Image.open(BytesIO(image_bytes))
        image = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
        
        # 提取数独网格
        grid_image = extract_sudoku_grid(image)
        
        # 分割为单元格
        cells = split_grid_to_cells(grid_image)
        
        # 识别数字
        sudoku_grid = []
        config = app.config['APP_CONFIG']
        ocr_engine = config.get('ocrEngine', 'opencv')
        
        for i in range(9):
            row = []
            for j in range(9):
                cell = cells[i][j]
                if ocr_engine == 'tesseract':
                    digit = recognize_digit_with_tesseract(cell, config)
                else:
                    digit = recognize_digit_with_opencv(cell)
                row.append(digit)
            sudoku_grid.append(row)
        
        return jsonify({
            'success': True,
            'grid': sudoku_grid,
            'message': 'OCR处理完成'
        })
        
    except Exception as e:
        print(f"OCR处理出错: {e}")  # 添加错误日志
        return jsonify({'success': False, 'error': str(e)}), 500

if __name__ == '__main__':
    # 确保配置文件存在
    if not os.path.exists(CONFIG_FILE):
        save_config(DEFAULT_CONFIG)
    
    app.run(host='0.0.0.0', port=5000, debug=True)